import os
from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QLabel, 
                             QComboBox, QGroupBox, QCheckBox, QRadioButton)
from PyQt6.QtCore import pyqtSignal

class Waifu2xCaffeParams(QWidget):
    """Waifu2x-Caffe参数面板"""
    
    # 参数变化信号
    params_changed = pyqtSignal(dict)
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.model_path = os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            'models',
            'waifu2x-caffe'
        )
        self.current_params = {
            'model_type': 'anime_style_art_rgb',  # 默认使用动漫风格模型
            'noise_level': 1,                    # 默认降噪等级1
            'scale': 2,                          # 默认放大倍数2x
            'tta_mode': False,                   # 默认关闭TTA模式
            'processor': 'gpu',                  # 默认使用GPU处理
        }
        
        self.init_ui()
    
    def init_ui(self):
        """初始化UI"""
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(10, 10, 10, 10)
        
        # 模型设置组
        model_group = QGroupBox("模型设置")
        model_layout = QVBoxLayout()
        
        # 模型类型
        model_type_layout = QHBoxLayout()
        model_type_layout.addWidget(QLabel("模型类型:"))
        self.model_type_combo = QComboBox()
        
        # 添加所有Waifu2x-Caffe模型类型
        self.model_types = [
            ("anime_style_art_rgb", "动漫风格 (RGB通道)"),
            ("anime_style_art", "动漫风格 (Y通道)"),
            ("photo", "照片风格"),
            ("upconv_7_anime_style_art_rgb", "UpResNet10 动漫风格 (RGB通道)"),
            ("upconv_7_photo", "UpResNet10 照片风格"),
            ("cunet", "CUnet模型")
        ]
        
        for model_id, model_name in self.model_types:
            self.model_type_combo.addItem(model_name, model_id)
            
        self.model_type_combo.currentIndexChanged.connect(self.on_model_changed)
        model_type_layout.addWidget(self.model_type_combo)
        model_layout.addLayout(model_type_layout)
        
        # 降噪等级
        noise_level_layout = QHBoxLayout()
        noise_level_layout.addWidget(QLabel("降噪等级:"))
        self.noise_level_combo = QComboBox()
        self.noise_level_combo.addItems(["0 (无降噪)", "1 (低)", "2 (中)", "3 (高)"])
        self.noise_level_combo.setCurrentIndex(1)  # 默认选择等级1
        self.noise_level_combo.currentIndexChanged.connect(self.on_params_changed)
        noise_level_layout.addWidget(self.noise_level_combo)
        model_layout.addLayout(noise_level_layout)
        
        # 放大倍数
        scale_layout = QHBoxLayout()
        scale_layout.addWidget(QLabel("放大倍数:"))
        self.scale_combo = QComboBox()
        self.scale_combo.addItems(["1x", "2x", "4x", "8x"])
        self.scale_combo.setCurrentIndex(1)  # 默认选择2x
        self.scale_combo.currentIndexChanged.connect(self.on_params_changed)
        scale_layout.addWidget(self.scale_combo)
        model_layout.addLayout(scale_layout)
        
        # 处理器选择
        processor_group = QGroupBox("处理器选择")
        processor_layout = QHBoxLayout()
        
        self.gpu_radio = QRadioButton("GPU (速度更快)")
        self.cpu_radio = QRadioButton("CPU (兼容性更好)")
        self.gpu_radio.setChecked(True)  # 默认选择GPU
        self.gpu_radio.toggled.connect(self.on_processor_changed)
        
        processor_layout.addWidget(self.gpu_radio)
        processor_layout.addWidget(self.cpu_radio)
        processor_group.setLayout(processor_layout)
        model_layout.addWidget(processor_group)
        
        # TTA模式
        self.tta_checkbox = QCheckBox("启用TTA模式 (提高质量但处理更慢)")
        self.tta_checkbox.stateChanged.connect(self.on_params_changed)
        model_layout.addWidget(self.tta_checkbox)
        
        model_group.setLayout(model_layout)
        main_layout.addWidget(model_group)
        
        # 添加弹性空间
        main_layout.addStretch(1)
        
        self.setLayout(main_layout)
    
    def on_model_changed(self, index):
        """模型类型改变时的处理函数"""
        self.current_params['model_type'] = self.model_type_combo.currentData()
        self.params_changed.emit(self.current_params)
    
    def on_processor_changed(self):
        """处理器选择改变时的处理函数"""
        self.current_params['processor'] = 'gpu' if self.gpu_radio.isChecked() else 'cpu'
        self.params_changed.emit(self.current_params)
    
    def on_params_changed(self):
        """参数变化时的处理函数"""
        # 更新参数
        self.current_params['noise_level'] = self.noise_level_combo.currentIndex()
        self.current_params['scale'] = [1, 2, 4, 8][self.scale_combo.currentIndex()]
        self.current_params['tta_mode'] = self.tta_checkbox.isChecked()
        
        # 发送参数变化信号
        self.params_changed.emit(self.current_params)
    
    def get_params(self):
        """获取当前参数"""
        return self.current_params
x